import pika
import json
import traceback
from process import predict

from config import RABBITMQ_CONFIG
print ('RABBITMQ_CONFIG = ', RABBITMQ_CONFIG)

from rabbitmq_utils import QUEUE_NAME, MULTITASK_QUEUE_NAME

credentials = pika.PlainCredentials(RABBITMQ_CONFIG['username'], RABBITMQ_CONFIG['password'])
parameters = pika.ConnectionParameters(
    host=RABBITMQ_CONFIG['host'],
    port=RABBITMQ_CONFIG['port'],
    credentials=credentials
)


# RabbitMQ 消费线程
def rabbitmq_consumer(socketio):
    # 单任务的消费者
    connection = pika.BlockingConnection(parameters)
    channel = connection.channel()
    channel.queue_declare(queue=QUEUE_NAME, durable=True)

    def callback(ch, method, properties, body):
        try:
            msg_str = body.decode('utf-8')
            result_str = predict(msg_str)
            result = json.loads(result_str)
            # 如果predict返回status=error，转为index=-1
            if result.get('status') == 'error':
                push_data = {'index': -1, 'errMsg': result.get('message', '未知错误')}
            else:
                push_data = result
            socketio.emit('message', push_data, namespace='/ws')
        except Exception as e:
            traceback.print_exc()
            err_data = {'index': -1, 'errMsg': str(e)}
            socketio.emit('message', err_data, namespace='/ws')

    channel.basic_consume(queue=QUEUE_NAME, on_message_callback=callback, auto_ack=True)
    print('RabbitMQ ws_consumer started.')
    channel.start_consuming()

def rabbitmq_multitask_consumer(socketio):    
    # 多任务的消费者
    connection = pika.BlockingConnection(parameters)
    channel = connection.channel()
    channel.queue_declare(queue=MULTITASK_QUEUE_NAME, durable=True)
    
    def callback(ch, method, properties, body):
        try:
            msg_str = body.decode('utf-8')
            result = json.loads(msg_str)
            # 如果predict返回status=error，转为index=-1
            if result.get('status') == 'error':
                push_data = {'index': -1, 'errMsg': result.get('message', '未知错误')}
            else:
                push_data = result
            socketio.emit('message', push_data, namespace='/multitask')
        except Exception as e:
            traceback.print_exc()
            err_data = {'index': -1, 'errMsg': str(e)}
            socketio.emit('message', err_data, namespace='/multitask')

    channel.basic_consume(queue=MULTITASK_QUEUE_NAME, on_message_callback=callback, auto_ack=True)
    print('RabbitMQ ws_multitask_consumer started.')
    channel.start_consuming()
